import torch
from torch.nn import Module, Linear, Softmax
from sharedinception import SharedInception
import torch.nn.functional as F
from fftblock import FFTBlock


class CnnSeasonBlock(Module):
    def __init__(self, topk:int, model_dim:int):
        super(CnnSeasonBlock, self).__init__()
        self.topk = topk
        self.fftblock = FFTBlock(self.topk)
        self.inception = SharedInception(model_dim)
        self.softmax = Softmax(dim=-1)
        

    def forward(self, season:torch.Tensor):
        # 对season进行傅里叶分解，得到周期和振幅
        amp, freq = self.fftblock.decomp(season)
        # print("amp and freq(shape) " + str(amp.shape) + " " + str(freq.shape))

        # 进行周期裁剪
        B, T, N = season.shape
        # print("B is " + str(B) + " T is " + str(T) + " N is " + str(N))
        segmented_data = []
        for i in range(freq.shape[-1]):
            f = freq[i].item()
            num_segments = (T  + f - 1) // f
            total_length = num_segments * f
            padded = F.pad(season, (0, 0, 0, total_length - T))
            # print("padded shape:" + str(padded.shape))
            segment = padded.view(B, num_segments, f, N)
            # print("segment shape:" + str(segment.shape))
            segmented_data.append(segment.permute(0, 3, 1, 2).contiguous())

        # 将segmented_data交给inception_block
        inception_outputs = []
        for segment in segmented_data:
            # print("segment shape:" + str(segment.shape))
            output = self.inception(segment)
            inception_outputs.append(output)
        
        # reshape back
        for i in range(len(inception_outputs)):
            inception_outputs[i] = inception_outputs[i].permute(0, 2, 3, 1).contiguous().reshape(B, -1, N)
            inception_outputs[i] = inception_outputs[i][:, :T, :]

        # 聚合结果
        res = torch.stack(inception_outputs, dim=-1)
        # print("amp", end=" ")
        # print(amp)
        amp = self.softmax(amp).unsqueeze(1).unsqueeze(1).repeat(1, T, N, 1)
        res = torch.sum(res * amp, dim=-1)

        # 残差连接
        return res + season


if __name__ == "__main__":
    B = 2
    T = 13
    N = 4
    x = torch.ones(B, T, N)
    block = CnnSeasonBlock(topk=3, model_dim=4)
    res = block(x)